# 数据加载和训练
class NovelDataLoader:
    def __init__(self, base_path):
        self.base_path = base_path
    
    def load_novels(self):
        """加载所有小说数据"""
        texts = []
        labels = []
        
        for category in os.listdir(self.base_path):
            category_path = os.path.join(self.base_path, category)
            if os.path.isdir(category_path):
                for file in os.listdir(category_path):
                    if file.endswith('.txt'):
                        file_path = os.path.join(category_path, file)
                        try:
                            with open(file_path, 'r', encoding='utf-8') as f:
                                content = f.read()
                                # 只取前10000字符进行训练，避免内存问题
                                texts.append(content[:10000])
                                labels.append(category)
                        except Exception as e:
                            print(f"读取文件 {file_path} 时出错: {e}")
        
        return texts, labels

def main():
    # 配置路径
    data_path = "novels"  # 你的小说数据路径
    model_save_path = "novel_classifier.pkl"
    
    # 加载数据
    print("加载数据...")
    loader = NovelDataLoader(data_path)
    texts, labels = loader.load_novels()
    
    print(f"总共加载 {len(texts)} 篇小说，分类: {set(labels)}")
    
    # 训练模型
    print("开始训练模型...")
    classifier = EnsembleNovelClassifier()
    X_test, y_test = classifier.train(texts, labels)
    
    # 测试集成模型
    test_texts = [texts[i] for i in range(min(5, len(texts)))]
    test_labels = [labels[i] for i in range(min(5, len(labels)))]
    
    predictions, probabilities = classifier.predict_ensemble(test_texts)
    
    print("\n测试结果:")
    for i, (true, pred, prob) in enumerate(zip(test_labels, predictions, probabilities)):
        print(f"样本 {i+1}: 真实={true}, 预测={pred}, 置信度={max(prob):.3f}")
    
    # 保存模型
    classifier.save_model(model_save_path)
    print(f"\n模型已保存到: {model_save_path}")

# 分类新小说的函数
def classify_new_novel(model_path, novel_path):
    """分类新的小说文件"""
    classifier = EnsembleNovelClassifier()
    classifier.load_model(model_path)
    
    with open(novel_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    prediction, probability = classifier.predict_ensemble([content])
    
    print(f"小说分类结果: {prediction[0]}")
    print(f"置信度: {max(probability[0]):.3f}")
    
    # 显示各类别概率
    categories = classifier.label_encoder.classes_
    for category, prob in zip(categories, probability[0]):
        print(f"  {category}: {prob:.3f}")
    
    return prediction[0]

if __name__ == "__main__":
    main()